#include "stream_puller_net.h"

namespace stream {

HBStreamPullerNET::HBStreamPullerNET(StreamPullerInfo* spinfo, LogInfo* log_info):
  HBStreamPullerBase(spinfo, log_info) {

}

HBStreamPullerNET::~HBStreamPullerNET() {

}

int HBStreamPullerNET::get_frame(stream::ImageBlob& buffer) {
  int ret = 0;
  queue_lock.lock();
  if (queue_data.empty()) { queue_lock.unlock(); return -1; }
  buffer = queue_data.front();
  queue_data.pop();
  queue_lock.unlock();
  return 0;
}

int HBStreamPullerNET::init_params() {
  int ret = 0;
  char errbuf[64];

  av_dict_set(&optionsDict, "stimeout", "3000000", 0);           // if don't setting this property，av_read_frame will run as block mode (ms)
  av_dict_set(&optionsDict, "bufsize", "1024000", 0);            // buffer size
  av_dict_set(&optionsDict, "rtsp_transport", puller_ifo->transferType.c_str(), 0); // transfer type,udp will faster but may lost some packet,tcp slower but stable
  if ((ret = avformat_open_input(&ifmt_ctx, puller_ifo->streamAddress.c_str(), 0, &optionsDict)) < 0) { // Open the input file for reading.
    spdlog::get("logger")->info("Could not open input file '{}' (error '{}')\n", puller_ifo->streamAddress, av_make_error_string(errbuf, sizeof(errbuf), ret));
    return -103;
  }
  spdlog::get("logger")->info("avformat_open_input ok!\n");

  if ((ret = avformat_find_stream_info(ifmt_ctx, nullptr)) < 0) { // Get information on the input file (number of streams etc.).
    spdlog::get("logger")->info("Could not open find stream info (error '%s')\n", av_make_error_string(errbuf, sizeof(errbuf), ret));
    return -102;
  } 
  spdlog::get("logger")->info("avformat_find_stream_info ok!\n");

  // // dump information
  // for (unsigned int i = 0; i < ifmt_ctx->nb_streams; i++) { av_dump_format(ifmt_ctx, i, puller_ifo->streamAddress.c_str(), 0); }
  // printf("av_dump_format ok!\n");

  // find video stream index
  for (unsigned int i = 0; i < ifmt_ctx->nb_streams; i++) { 
    AVStream *st = ifmt_ctx->streams[i];
    if (st->codec->codec_type == AVMEDIA_TYPE_AUDIO) { audio_st_index = i; }
    else if (st->codec->codec_type == AVMEDIA_TYPE_VIDEO) { video_st_index = i; }
  }

  if (-1 == video_st_index) {
    printf("No H.264 video stream in the input file\n");
    return -100;
  }

  // getting stream msg
  bool nRestart = false;
  if (!nRestart) { 
    pVst = ifmt_ctx->streams[video_st_index];
    pVideoCodecCtx = pVst->codec;
    if (pVideoCodecCtx->codec_id == AV_CODEC_ID_H264) { encode_type = PT_H264;}
    else if (pVideoCodecCtx->codec_id == AV_CODEC_ID_H265) { encode_type  = PT_H265;}

    // 无需ffmpeg解码
    // pVideoCodec = avcodec_find_decoder(pVideoCodecCtx->codec_id);
    // if (pVideoCodec == nullptr) {return;}
    // if (avcodec_open2(pVideoCodecCtx, pVideoCodec, nullptr) < 0) {return;}
  }
  if (encode_type == PT_PCMU) {
    printf("[ERROR] sp_encoder error\n");
    return -101;
  }
  /*Begin decoder and display*/
  // getting stream height and width
  puller_ifo->oriWidth = pVideoCodecCtx->width;
  puller_ifo->oriHeight = pVideoCodecCtx->height;
  spdlog::get("logger")->info("decodeChn:{},oriWidth:{},oriHeight:{}\n", puller_ifo->vdecChn, puller_ifo->oriWidth, puller_ifo->oriHeight);

  stop_flag = common::CameraOpened;
  return 0;
}




int HBStreamPullerNET::start_pull() {
  int ret = 0;

  start_pull_imp();

  release_ffmpeg();
  return 0;

}

int HBStreamPullerNET::start_pull_imp() {
  int ret = 0;

  // 初始化channel属性
  VDEC_CHN_ATTR_S vdecChnAttr;
  ret = sample_vdec_ChnAttr_init(&vdecChnAttr);
  if (ret != 0) { spdlog::get("logger")->info("ERROR. sample_venc_ChnAttr_init failded: {}\n", ret);}
  // 创建channel
  ret = HB_VDEC_CreateChn(puller_ifo->vdecChn, &vdecChnAttr);
  if (ret != 0) { spdlog::get("logger")->info("ERROR. HB_VDEC_CreateChn {} failed, {}.\n", puller_ifo->vdecChn, ret);}
  // 设置channel属性
  ret = HB_VDEC_SetChnAttr(puller_ifo->vdecChn, &vdecChnAttr);  // config
  if (ret != 0) { spdlog::get("logger")->info("ERROR. HB_VDEC_SetChnAttr failed {}\n", ret);}

  ret = HB_VDEC_StartRecvStream(puller_ifo->vdecChn);
  if (ret != 0) { spdlog::get("logger")->info("ERROR. HB_VDEC_StartRecvStream failed {}\n", ret);}


  std::thread(&HBStreamPullerNET::start_pull_listener, this).detach();
  std::this_thread::sleep_for(std::chrono::milliseconds(3000));

  int size_y = puller_ifo->oriWidth * puller_ifo->oriHeight;
  int size_uv = size_y / 2;
  stream::ImageBlob buffer(stream::ImageBlobMode_YUV, size_y+size_uv);
  buffer.ori_im_shape = {puller_ifo->oriHeight, puller_ifo->oriWidth};
  buffer.yuv_size = size_y+size_uv;

  int frame_cnt = 0;
  VIDEO_FRAME_S stFrameInfo;
  // int64_t timepoint = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();;
  int idx = 0;
  float vdec_fps = 0;
  int64_t cur_time = 0;
  int vdec_cnt = 0;
  int64_t start_time = 0, end_time = 0;
  stop_flag2 = common::CameraGrabing;
  while (stop_flag2 == common::CameraGrabing) {
    std::this_thread::sleep_for(std::chrono::milliseconds(3));
    ret = HB_VDEC_GetFrame(puller_ifo->vdecChn, &stFrameInfo, 1000);
    // printf("HB_VDEC_GetFrame:%d\n", ret);
    if (ret != 0) { printf("HB_VDEC_GetFrame failed:%d\n", ret); continue; }
    vdec_cnt++;

    cur_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    std::copy(stFrameInfo.stVFrame.vir_ptr[0], stFrameInfo.stVFrame.vir_ptr[0] + size_y, buffer.yuv.begin());
    std::copy(stFrameInfo.stVFrame.vir_ptr[1], stFrameInfo.stVFrame.vir_ptr[1] + size_uv, buffer.yuv.begin()+size_y);
    if (idx % 300 == 0) {printf("vdec %d, %dx%d, %lld\n", idx++, stFrameInfo.stVFrame.width, stFrameInfo.stVFrame.height, cur_time);}

    HB_VDEC_ReleaseFrame(puller_ifo->vdecChn, &stFrameInfo);

    queue_lock.lock();
    if (queue_data.size() >= buffer_cnt) {queue_data.pop();}
    queue_data.push(buffer);
    queue_lock.unlock();

    // 输出fps
    end_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    if (end_time - start_time > 5 * 1e3) {
      vdec_fps = 1.0 / ((end_time - start_time) / (vdec_cnt * 1e3));
      printf("vdec_fps c:%d : %f, vdec_cnt:%d\n", puller_ifo->vdecChn, vdec_fps, vdec_cnt);
      spdlog::get("logger")->info("RKStreamPullerNET::start_pull. vdec_fps c:{} : {}, vdec_cnt:{}", puller_ifo->channelId, vdec_fps, vdec_cnt);
      vdec_cnt = 0;
      start_time = end_time;
    }
  }
  
  // free
  while (!queue_data.empty()) { queue_data.pop(); }
  stop_flag2 = common::CameraClosed;
  return 0;

}

int HBStreamPullerNET::start_pull_listener() {
  int ret = 0;
  // stop_flag = common::CameraGrabing;
  // while (stop_flag == common::CameraGrabing) {
  //   if (is_init == 0 && (ret = init_params()) != 0) {
  //     spdlog::get("logger")->info("ERROR init_params error ret:{}\n", ret);
  //     exit(0);
  //   }

    ret = start_pull_listener_imp();
  // }
  // stop_flag = common::CameraOpened;
  return ret;
}

int HBStreamPullerNET::start_pull_listener_imp() {
  int ret = 0;
  // char errbuf[64];

  AVPacket pkt;
  av_init_packet(&pkt);    // initialize packet.
  pkt.data = nullptr;
  pkt.size = 0;


  // 准备buffer
  int mmz_cnt = 3;
  uint64_t mmz_paddr[mmz_cnt];
  memset(mmz_paddr, 0, sizeof(mmz_paddr));
  char* mmz_vaddr[mmz_cnt];
  int mmz_size = puller_ifo->oriWidth * puller_ifo->oriHeight / 4;
  for (int i = 0; i < mmz_cnt; i++) {
    mmz_vaddr[i] = nullptr;
    ret = HB_SYS_Alloc(&(mmz_paddr[i]), (void **)&mmz_vaddr[i], mmz_size);
    // printf("mmzAlloc paddr = 0x%lx, vaddr = 0x%lx i = %d, ret=%d \n", (uint64_t)mmz_paddr[i], (uint64_t)mmz_vaddr[i],i, ret); 
    if (ret != 0) { spdlog::get("logger")->info("mmzAlloc paddr = 0x{}, vaddr = 0x{} i = {}, ret={} \n", (uint64_t)mmz_paddr[i], (uint64_t)mmz_vaddr[i],i, ret);}
  }
  VIDEO_STREAM_S pstStream;
  memset(&pstStream, 0, sizeof(VIDEO_STREAM_S));

  /*End*/
  uint8_t *seqHeader = nullptr;
  int seqHeaderSize = 0;
  int firstPacket = 1;
  int eos = 0;
  int bufSize = 0;
  float pull_fps;
  int pts_count = 0, fps_cnt = 0;
  int64_t start_time = 0, end_time = 0;
  int mmz_index = 0;
  VDEC_CHN_STATUS_S pstStatus;
  int av_error=0;
  stop_flag = common::CameraGrabing;
  while (stop_flag == common::CameraGrabing) {
    std::this_thread::sleep_for(milliseconds(3));

    ret = HB_VDEC_QueryStatus(puller_ifo->vdecChn, &pstStatus);
    if (ret != 0) { spdlog::get("logger")->info("ERROR. HB_VDEC_QueryStatus error:{}\n", ret); break;}
    if (pstStatus.cur_input_buf_cnt >= (uint32_t)mmz_cnt) {
        std::this_thread::sleep_for(milliseconds(5));
        continue;
    }
    // std::cout << "pstStatus.cur_input_buf_cnt:" << pstStatus.cur_input_buf_cnt << std::endl;

    ret = av_read_frame(ifmt_ctx, &pkt); // read frames
    if (ret < 0) {
      av_error++; 
      if (av_error%50==0) {spdlog::get("logger")->info("ERROR av_read_frame ret < 0, av_error:{}\n", av_error); }
      // 退出重新拉流
      if (av_error > 1000) { av_packet_unref(&pkt); pkt.size = 0; break; }
      continue; 
    } else {
      if (av_error > 0) { av_error--; }
    }

    // printf("Could not read frame ---(error '%s')\n", av_make_error_string(errbuf, sizeof(errbuf), ret)); 

    seqHeaderSize = 0;
    mmz_index = pts_count % mmz_cnt;

    if (pkt.stream_index != video_st_index) { continue; }
    if (firstPacket) {
      int retSize = 0;
      AVCodecParameters* codec = ifmt_ctx->streams[video_st_index]->codecpar;
      seqHeader = (uint8_t*)malloc(codec->extradata_size + 1024);
      if (seqHeader == nullptr) {
        spdlog::get("logger")->info("ERROR. Failed to mallock seqHeader\n");
        eos = 1;
        break;
      }
      memset((void*)seqHeader, 0x00, codec->extradata_size + 1024);

      seqHeaderSize = build_dec_seq_header(seqHeader, ifmt_ctx->streams[video_st_index], &retSize);
      if (seqHeaderSize < 0) {
        spdlog::get("logger")->info("ERROR. Failed to build seqHeader\n");
        eos = 1;
        break;
      }
      firstPacket = 0;
    }
    if (pkt.size <= mmz_size) {
      if (seqHeaderSize) {
        memcpy((void*)mmz_vaddr[mmz_index], (void*)seqHeader, seqHeaderSize);
        bufSize = seqHeaderSize;
      } else {
        memcpy((void*)mmz_vaddr[mmz_index], (void*)pkt.data, pkt.size);
        bufSize = pkt.size;
      }
      av_packet_unref(&pkt);
      pkt.size = 0;
    } else {
      spdlog::get("logger")->info("ERROR. error. mmz_size:{}, pkt.size:{}\n", mmz_size, pkt.size);
    }
    if (seqHeader) { free(seqHeader); seqHeader = nullptr; }

    // 
    pstStream.pstPack.phy_ptr = mmz_paddr[mmz_index];
    pstStream.pstPack.vir_ptr = mmz_vaddr[mmz_index];
    pstStream.pstPack.pts = pts_count++;
    pstStream.pstPack.src_idx = mmz_index;
    if (!eos) {
      pstStream.pstPack.size = bufSize;
      pstStream.pstPack.stream_end = HB_FALSE;
    } else {
      pstStream.pstPack.size = 0;
      pstStream.pstPack.stream_end = HB_TRUE;
    }
    // printf("[pstStream] pts:%ld, vir_ptr:%ld, size:%d\n",
    //         pstStream.pstPack.pts,
    //         (uint64_t)pstStream.pstPack.vir_ptr,
    //         pstStream.pstPack.size);

    if (fps_cnt == 0) { start_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();}
    fps_cnt++;
    
    // printf("feed raw data\n");
    ret = HB_VDEC_SendStream(puller_ifo->vdecChn, &pstStream, 3000);
    // printf("feed raw data ret:%d\n", ret);
    if (ret != 0) { spdlog::get("logger")->info("ERROR. HB_VDEC_SendStream failed\n"); }

    // 输出fps
    end_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    if (end_time - start_time > 5 * 1e3) {
      pull_fps = 1.0 / ((end_time - start_time) / (fps_cnt * 1e3));
      if (log_ifo->log_level_2) spdlog::get("logger")->info("chn:{}, pull_fps:{}", puller_ifo->vdecChn, pull_fps);
      printf("pull_fps c:%d : %f\n", puller_ifo->vdecChn, pull_fps);
      fps_cnt = 0;
    }
  }


  HB_VDEC_DestroyChn(puller_ifo->vdecChn);
  for (int i = 0; i < mmz_cnt; i++) {
    ret = HB_SYS_Free(mmz_paddr[i], mmz_vaddr[i]);
    // printf("mmzFree paddr = 0x%lx, vaddr = 0x%lx i = %d, ret=%d\n", (uint64_t)mmz_paddr[i], (uint64_t)mmz_vaddr[i], i, ret);
    if (ret != 0) { spdlog::get("logger")->info("mmzFree paddr = 0x{}, vaddr = 0x{} i = {}, ret={}\n", (uint64_t)mmz_paddr[i], (uint64_t)mmz_vaddr[i], i, ret);}
  }
  release_ffmpeg();

  stop_flag = common::CameraClosed;
  return 0;
}

int HBStreamPullerNET::build_dec_seq_header(uint8_t* pbHeader, const AVStream* st, int* sizelength) {
  AVCodecParameters* avc = st->codecpar;

  uint8_t* pbMetaData = avc->extradata;
  int nMetaData = avc->extradata_size;
  uint8_t* p =    pbMetaData;
  uint8_t *a =    p + 4 - ((long) p & 3);
  uint8_t* t =    pbHeader;
  int         size;
  int         sps, pps, i, nal;

  size = 0;
  *sizelength = 4;  // default size length(in bytes) = 4
  if (encode_type == PT_H264) {
    if (nMetaData > 1 && pbMetaData && pbMetaData[0] == 0x01) {
      // check mov/mo4 file format stream
      p += 4;
      *sizelength = (*p++ & 0x3) + 1;
      sps = (*p & 0x1f);  // Number of sps
      p++;
      for (i = 0; i < sps; i++) {
        nal = (*p << 8) + *(p + 1) + 2;
        SET_BYTE(t, 0x00);
        SET_BYTE(t, 0x00);
        SET_BYTE(t, 0x00);
        SET_BYTE(t, 0x01);
        SET_BUFFER(t, p+2, nal-2);
        p += nal;
        size += (nal - 2 + 4);  // 4 => length of start code to be inserted
      }

      pps = *(p++);  // number of pps
      for (i = 0; i < pps; i++)
      {
        nal = (*p << 8) + *(p + 1) + 2;
        SET_BYTE(t, 0x00);
        SET_BYTE(t, 0x00);
        SET_BYTE(t, 0x00);
        SET_BYTE(t, 0x01);
        SET_BUFFER(t, p+2, nal-2);
        p += nal;
        size += (nal - 2 + 4);  // 4 => length of start code to be inserted
      }
    } else if(nMetaData > 3) {
      size = -1;  // return to meaning of invalid stream data;
      for (; p < a; p++) {
        if (p[0] == 0 && p[1] == 0 && p[2] == 1)  {
          // find startcode
          size = avc->extradata_size;
          if (!pbHeader || !pbMetaData)
            return 0;
          SET_BUFFER(pbHeader, pbMetaData, size);
          break;
        }
      }
    }
  } else if (encode_type == PT_H265) {
    if (nMetaData > 1 && pbMetaData && pbMetaData[0] == 0x01) {
      static const int8_t nalu_header[4] = { 0, 0, 0, 1 };
      int numOfArrays = 0;
      uint16_t numNalus = 0;
      uint16_t nalUnitLength = 0;
      uint32_t offset = 0;

      p += 21;
      *sizelength = (*p++ & 0x3) + 1;
      numOfArrays = *p++;

      while(numOfArrays--) {
        p++;   // NAL type
        numNalus = (*p << 8) + *(p + 1);
        p+=2;
        for(i = 0;i < numNalus;i++)
        {
          nalUnitLength = (*p << 8) + *(p + 1);
          p+=2;
          // if(i == 0)
          {
            memcpy(pbHeader + offset, nalu_header, 4);
            offset += 4;
            memcpy(pbHeader + offset, p, nalUnitLength);
            offset += nalUnitLength;
          }
          p += nalUnitLength;
        }
      }

      size = offset;
    } else if(nMetaData > 3) {
      size = -1;  // return to meaning of invalid stream data;

      for (; p < a; p++)
      {
        if (p[0] == 0 && p[1] == 0 && p[2] == 1)  // find startcode
        {
          size = avc->extradata_size;
          if (!pbHeader || !pbMetaData)
            return 0;
          SET_BUFFER(pbHeader, pbMetaData, size);
          break;
        }
      }
    }
  } else {
    SET_BUFFER(pbHeader, pbMetaData, nMetaData);
    size = nMetaData;
  }

  return size;
}

void HBStreamPullerNET::release_ffmpeg() {
  avcodec_close(pVideoCodecCtx);
  avformat_close_input(&ifmt_ctx);
  // if (pFrame != nullptr) av_free(pFrame);
  ifmt_ctx = nullptr;

  delete optionsDict;
  optionsDict = nullptr;
}

}